package com.csthink.limiter.aop;

import com.csthink.limiter.condition.LimitAspectCondition;
import com.csthink.limiter.config.LimiterProperties;
import com.csthink.limiter.manager.Limiter;
import com.csthink.limiter.exception.LimiterException;
import com.csthink.limiter.manager.LimiterManager;
import lombok.Setter;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Conditional;
import org.springframework.context.annotation.EnableAspectJAutoProxy;
import org.springframework.util.StringUtils;

import java.lang.reflect.Method;

/**
 * 通过 AOP 切面拦截限流注解
 *
 * @author <a href="mailto:security.2009@live.cn">Mars</a>
 * @since 2023-09-08
 */
@Aspect
@EnableAspectJAutoProxy(proxyTargetClass = true) // 使用CGLIB代理
@Conditional(LimitAspectCondition.class)
public class LimitAspect {

    @Setter(onMethod_ = @Autowired)
    private LimiterManager limiterManager;

    private final LimiterProperties limiterProperties;

    public LimitAspect(LimiterProperties limiterProperties) {
        this.limiterProperties = limiterProperties;
    }

    @Pointcut("@annotation(com.csthink.limiter.aop.Limit)")
    private void check() {

    }

    @Before("check()")
    public void before(JoinPoint joinPoint) {

        if (!limiterProperties.isEnabled()) {
            return;
        }

        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        // 获取方法签名上的 Limit 注解,如果存在则说明需要限流
        Limit limit = method.getAnnotation(Limit.class);
        if (limit == null) {
            return;
        }

        String key = limit.key();
        if (!StringUtils.hasText(key)) {
            throw new LimiterException("Limiter key cannot be null");
        }

        if (limit.permitsPerSecond() <= 0 || limit.timeout() <= 0) {
            throw new LimiterException("Limiter param invalid");
        }

        Limiter limiter = Limiter.builder()
                .key(key)
                .permitsPerSecond(limit.permitsPerSecond())
                .timeout(limit.timeout())
                .build();

        if (!limiterManager.tryAccess(limiter)) {
            throw new LimiterException("Current Server is busy, please try again later!");
        }
    }
}
